package com.zhuangjie.mvc.future;

import io.seata.core.context.RootContext;
import org.slf4j.MDC;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.ForkJoinPool;

/**
 * 上下文感知执行器
 *
 * @author zhuangjie
 * @date 2024/01/03
 */
public class ContextAwareExecutor implements Executor {
    static final class ThreadPerTaskExecutor implements Executor {
        public void execute(Runnable r) { new Thread(r).start(); }
    }
    private final Executor executor;

    public ContextAwareExecutor() {
        executor = ForkJoinPool.getCommonPoolParallelism() > 1 ?
                ForkJoinPool.commonPool() : new ContextAwareExecutor.ThreadPerTaskExecutor();
    }

    public ContextAwareExecutor(Executor executor) {
        this.executor = executor;
    }

    @Override
    public void execute(Runnable command) {
        // TOKEN: 将父线程的上下文捕获并传递给子线程
        RequestAttributes context = RequestContextHolder.currentRequestAttributes();
        // SEATA: 分布式事务id
        String xid = RootContext.getXID();
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();

        executor.execute(() -> {
            try {
                // TOKEN: 向下传入token
                RequestContextHolder.setRequestAttributes(context);
                // SEATA: 向下传递seata分布式事务id，使得异步操作也处于分布式事务内
                if(xid != null) RootContext.bind(xid);
                // Restore the MDC context in the new thread
                if (mdcContext != null) {
                    MDC.setContextMap(mdcContext);
                }
                command.run();

            } finally {
                RequestContextHolder.resetRequestAttributes();
                RootContext.unbind();
                MDC.clear();
            }
        });
    }
}
